import exrex, re, random, string
from datetime import datetime, timedelta
import pandas as pd
import json
import os
from typing import List, Any, Union, Dict


class RuleAnalysis(object):

    _rule = None

    _choice_cache = {}  # 缓存外部 choice 数据，避免重复读取

    # 🆕 按文件路径缓存当前选中的行 {file_path: row_dict}
    _current_choice_row = {}

    _current_time = None

    def __init__(self) -> None:
        pass

    def _setrule(self, rule_reg):
        self._rule = rule_reg
        return self

    def _extract_json_path(self, data, path: str) -> List[Any]:
        """简易 JSON Path 提取（支持 data[*].field）"""
        parts = path.split(".")
        current = data

        for part in parts:
            if part.endswith("[*]"):
                key = part[:-3]
                if isinstance(current, dict) and key in current:
                    current = current[key]
                else:
                    return []
                if not isinstance(current, list):
                    return []
                # 继续下一层
                results = []
                for item in current:
                    if isinstance(item, dict) and len(parts) > parts.index(part) + 1:
                        sub_path = ".".join(parts[parts.index(part) + 1 :])
                        sub_values = self._extract_json_path(item, sub_path)
                        results.extend(sub_values)
                    elif isinstance(item, dict):
                        # 最后一层
                        next_key = (
                            parts[parts.index(part) + 1]
                            if len(parts) > parts.index(part) + 1
                            else None
                        )
                        if next_key and next_key in item:
                            results.append(item[next_key])
                        else:
                            results.append(item)
                    else:
                        results.append(item)
                return results
            else:
                if isinstance(current, dict) and part in current:
                    current = current[part]
                else:
                    return []

        if isinstance(current, list):
            return current
        else:
            return [current]

    def _flatten_json(self, obj) -> List[Any]:
        """递归扁平化 JSON 所有叶子值"""
        results = []
        if isinstance(obj, dict):
            for v in obj.values():
                results.extend(self._flatten_json(v))
        elif isinstance(obj, list):
            for item in obj:
                results.extend(self._flatten_json(item))
        else:
            results.append(obj)
        return results

    def _extract_excel_column(
        self, file_path: str, field_spec: str = None
    ) -> List[Any]:
        """从 Excel 提取某一列数据"""
        if not field_spec:
            # 默认读取第一个 Sheet 第一列
            df = pd.read_excel(file_path, sheet_name=0)
            col = df.iloc[:, 0]  # 第一列
        elif "!" in field_spec:
            sheet_name, col_name = field_spec.split("!", 1)
            df = pd.read_excel(file_path, sheet_name=sheet_name)
            col = df[col_name]
        else:
            # 只给列名，默认第一个 Sheet
            df = pd.read_excel(file_path, sheet_name=0)
            col = df[field_spec]

        return col.dropna().astype(str).tolist()  # 转为字符串列表，去掉空值

    def _load_choice_file(self, file_path: str):
        """统一加载文件到缓存"""
        # 检查文件是否存在
        if not os.path.exists(file_path):
            raise FileNotFoundError(f"文件不存在: {file_path}")

        ext = os.path.splitext(file_path)[1].lower()
        try:
            if ext == ".json":
                with open(file_path, "r", encoding="utf-8") as f:
                    data = json.load(f)
                if isinstance(data, list):
                    self._choice_cache[file_path] = data
                else:
                    self._choice_cache[file_path] = self._flatten_json(data)
            elif ext in [".xlsx", ".xls"]:
                self._choice_cache[file_path] = pd.read_excel(file_path, sheet_name=0)
            elif ext == ".csv":
                self._choice_cache[file_path] = pd.read_csv(file_path)
            else:
                raise ValueError(f"不支持的文件类型: {ext}")
        except Exception as e:
            raise ValueError(f"加载文件失败: {file_path}, 错误: {str(e)}")

    def format_choice_ref(self, match):
        """
        支持：
          - choice_ref@字段名          → 引用最近一次 choice_row 选中的文件中的字段
          - choice_ref@文件名!字段名   → 明确引用指定文件缓存的行
        """
        ref_spec = match.group(1).strip()

        if "!" in ref_spec:
            file_part, field_name = ref_spec.split("!", 1)
            current_row = self._current_choice_row.get(file_part, {})
            if not current_row:
                return f"NO_ROW_CACHED_FOR_FILE:{file_part}"
        else:
            field_name = ref_spec
            if not self._current_choice_row:
                return "NO_ROW_CACHED"
            # 取最后被设置的文件（按插入顺序，取最后一个 key）
            last_file = next(reversed(self._current_choice_row))
            current_row = self._current_choice_row.get(last_file, {})
            if not current_row:
                return "NO_ROW_CACHED"

        if field_name in current_row:
            return str(current_row[field_name])
        else:
            return f"FIELD_NOT_FOUND:{field_name}"

    def format_choice(self, match):
        """
        支持：
          - choice@["A", "B", "C"]
          - choice@file.xlsx!ColumnName
          - choice_row@file.xlsx!KeyColumn → 选中整行，按文件缓存供后续字段使用
        """
        full_match = match.group(0)
        param = match.group(1).strip()

        if param.startswith("[") and param.endswith("]"):
            try:
                choices = json.loads(param)
                if not isinstance(choices, list):
                    raise ValueError
                selected = random.choice(choices)
                return str(selected)
            except:
                return "INVALID_CHOICE_ARRAY"

        # 支持 choice_row@ 语法（选中整行并缓存）
        is_row_mode = "choice_row@" in full_match

        # 解析文件路径和字段
        if "!" in param:
            file_part, field_part = param.split("!", 1)
        else:
            file_part, field_part = param, None

        # 读取数据
        if file_part not in self._choice_cache:
            self._load_choice_file(file_part)

        df_or_list = self._choice_cache[file_part]

        if isinstance(df_or_list, list):
            # JSON 数据（列表）
            if field_part:
                # 提取指定字段的值
                values = [
                    item.get(field_part)
                    for item in df_or_list
                    if isinstance(item, dict)
                ]
                values = [v for v in values if v not in (None, "")]
            else:
                values = df_or_list
            if not values:
                return "NO_DATA"
            selected_value = random.choice(values)

            # 如果是 row 模式，缓存整行（按文件）
            if is_row_mode and field_part:
                # 找到选中的那个 item
                candidates = [
                    item
                    for item in df_or_list
                    if isinstance(item, dict) and item.get(field_part) == selected_value
                ]
                if candidates:
                    self._current_choice_row[file_part] = candidates[0]
                else:
                    # 清除该文件缓存
                    if file_part in self._current_choice_row:
                        del self._current_choice_row[file_part]
            elif not is_row_mode:
                # 非 row 模式：清除该文件的行缓存（可选）
                if file_part in self._current_choice_row:
                    del self._current_choice_row[file_part]

            return str(selected_value)

        else:
            # DataFrame (Excel/CSV)
            if not field_part:
                field_part = df_or_list.columns[0]  # 默认第一列

            if field_part not in df_or_list.columns:
                return f"COLUMN_NOT_FOUND:{field_part}"

            # 随机选一行
            row = df_or_list.sample(1).iloc[0]
            selected_value = row[field_part]

            # 🆕 按文件路径缓存整行
            if is_row_mode:
                self._current_choice_row[file_part] = row.to_dict()
            else:
                # 非 row 模式：清除该文件的行缓存（可选）
                if file_part in self._current_choice_row:
                    del self._current_choice_row[file_part]

            return str(selected_value)

    def set_current_time(self, dt: datetime):
        """设置当前时间上下文"""
        self._current_time = dt

    def formatdate(self, exrex_str):
        now = self._current_time or datetime.now()
        date_format_map = {
            "yyyy": lambda: now.strftime("%Y"),
            "MM": lambda: now.strftime("%m"),
            "dd": lambda: now.strftime("%d"),
            "HH": lambda: now.strftime("%H"),
            "mm": lambda: now.strftime("%M"),
            "ss": lambda: now.strftime("%S"),
        }
        keys = date_format_map.keys()
        if any(key in exrex_str for key in keys):
            for placeholder, formatter in date_format_map.items():
                exrex_str = exrex_str.replace(placeholder, formatter())
        return exrex_str

    def random_string(self, count=20, letters=True, numbers=False):
        if count <= 0:
            count = 20

        char_pool = ""

        if letters:
            char_pool += string.ascii_letters  # 包含所有字母
        if numbers:
            char_pool += string.digits  # 包含数字

        if not char_pool:  # 如果没有字母和数字
            char_pool = "".join(
                chr(i) for i in range(0, 0x110000)
            )  # 包含所有 Unicode 字符

        buffer = []

        while count > 0:
            ch = random.choice(char_pool)  # 随机选择字符
            buffer.append(ch)
            count -= 1

        return "".join(buffer)

    # 定义替换函数
    def format_uid(self, match):
        # 提取匹配到的数字
        try:
            number = match.group(1)
        except:
            number = 20

        # 转换为整数，若未匹配到则使用默认值20
        value = int(number) if number else 20
        return self.random_string(value, False, True)  # 返回替换后的字符串

    # 定义替换函数
    def format_ipv4(self, match):
        # IPv4 地址的正则表达式
        ipv4_pattern = r"(?:25[0-5]|2[0-4][0-9]|[1-9][0-9]?|0)\.(?:25[0-5]|2[0-4][0-9]|[1-9][0-9]?|0)\.(?:25[0-5]|2[0-4][0-9]|[1-9][0-9]?|0)\.(?:25[0-5]|2[0-4][0-9]|[1-9][0-9]?|0)"
        return exrex.getone(ipv4_pattern)  # 返回替换后的字符串

    def format_ipv6(self, match):
        # IPv6 地址的正则表达式
        ipv6_pattern = r"([0-9A-F]{4}:){7}[0-9A-F]{4}"
        return exrex.getone(ipv6_pattern)  # 返回替换后的字符串

    # uid 替换
    def formatter(self, input_string, pattern, replacement):
        # 使用re.sub进行替换
        output_string = re.sub(pattern, replacement, input_string)
        return output_string

    def format_extra_str(self, exrex_str):
        # 定义正则表达式
        ipv4_search_pattern = r"ipv4"
        ipv6_search_pattern = r"ipv6"
        uid_search_pattern = r"uid(?:\((\d+)\))?"
        # 修复正则表达式：将 [^\s&&]+ 改为 [^\s&]+
        choice_search_pattern = r"choice(?:_row)?@([^&]+)"
        choice_ref_pattern = r"choice_ref@([^&]+)"
        extra_format_map = {
            "ipv4": lambda: self.formatter(
                exrex_str, ipv4_search_pattern, self.format_ipv4
            ),
            "ipv6": lambda: self.formatter(
                exrex_str, ipv6_search_pattern, self.format_ipv6
            ),
            "uid": lambda: self.formatter(
                exrex_str, uid_search_pattern, self.format_uid
            ),
            "choice": lambda: self.formatter(
                exrex_str, choice_search_pattern, self.format_choice
            ),
            "choice_ref": lambda: self.formatter(
                exrex_str, choice_ref_pattern, self.format_choice_ref
            ),
        }

        # 判断字符串中是否包含了date_format_map 的key
        keys = extra_format_map.keys()
        if any(key in exrex_str for key in keys):
            # 通过 date_format_map 替换exrex_str命中规则
            for placeholder, formatter in extra_format_map.items():
                if placeholder not in exrex_str:
                    continue

                exrex_str = formatter()

        return exrex_str

    def _mockstr(self):
        # 生成匹配给定正则表达式的第一个字符串
        exrex_str = "".join(exrex.getone(self._rule))
        # 格式化时间字符串
        exrex_str = self.formatdate(exrex_str)
        # 生成 其他类型参数
        exrex_str = self.format_extra_str(exrex_str)
        return exrex_str

    def _mocknum(self):
        exrex_num = None
        rule_str = self._rule.strip()

        # 新增：支持范围语法 num@120~3600
        if "~" in rule_str:
            try:
                low, high = map(int, rule_str.split("~", 1))
                if low > high:
                    low, high = high, low
                return str(random.randint(low, high))
            except (ValueError, TypeError):
                pass  # 如果解析失败，继续走原有逻辑

        # 原有逻辑：处理纯数字、正则等
        try:
            return str(float(rule_str)).rstrip("0").rstrip(".")
        except:
            try:
                length = int(rule_str)
            except:
                exrex_num = exrex.getone(rule_str)
                try:
                    float(exrex_num)
                except:
                    exrex_num = None
                    length = 1

        exrex_num = exrex.getone(f"\\d{{{length}}}") if exrex_num is None else exrex_num
        if "." not in exrex_num:
            return str(int(exrex_num))
        return exrex_num.rstrip("0").rstrip(".")

    def _mockdate(self):
        now = self._current_time or datetime.now()
        date_reg = r"(19|20)\d{2}-(0[1-9]|1[0-2])-(0[1-9]|[12]\d|3[01]) ([01]\d|2[0-3]):([0-5]\d):([0-5]\d)"
        if not self._rule:
            self._rule = date_reg
        exrex_str = exrex.getone(self._rule)

        # 正则表达式模式
        pattern = r"(?P<year>\d{4})-(?P<month>0[1-9]|1[0-2])-(?P<day>0[1-9]|[12][0-9]|3[01])(?: (?P<hour>[01][0-9]|2[0-3]):(?P<minute>[0-5][0-9])(?::(?P<second>[0-5][0-9])(?:\.(?P<micro>[0-9]{3}))?)?)?"

        # 初始化结果字典
        result = {
            "yyyy": None,
            "MM": None,
            "dd": None,
            "HH": None,
            "mm": None,
            "ss": None,
            "SSS": None,
        }
        matches = re.finditer(pattern, exrex_str)
        for match in matches:
            if match.group("year"):
                result["yyyy"] = match.group("year")
            if match.group("month"):
                result["MM"] = match.group("month")
            if match.group("day"):
                result["dd"] = match.group("day")
            if match.group("hour"):
                result["HH"] = match.group("hour")
            if match.group("minute"):
                result["mm"] = match.group("minute")
            if match.group("second"):
                result["ss"] = match.group("second")
            if match.group("micro"):
                result["SSS"] = match.group("micro")

        # 填充缺失的部分，使用当前时间
        result["yyyy"] = result["yyyy"] or str(now.year)
        result["MM"] = result["MM"] or f"{now.month:02}"
        result["dd"] = result["dd"] or f"{now.day:02}"
        result["HH"] = result["HH"] or f"{now.hour:02}"
        result["mm"] = result["mm"] or f"{now.minute:02}"
        result["ss"] = result["ss"] or f"{now.second:02}"
        result["SSS"] = result["SSS"] or f"{now.microsecond:03}"

        final_datetime = f"{result['yyyy']}-{result['MM']}-{result['dd']} {result['HH']}:{result['mm']}:{result['ss']}.{result['SSS']}"
        try:
            datetime.strptime(final_datetime, "%Y-%m-%d %H:%M:%S.%f")
            return final_datetime
        except:
            # 生成一个基于 current_time 的随机时间
            delta = timedelta(seconds=random.randint(0, 86400))  # 一天内随机
            random_time = now + delta
            return random_time.strftime("%Y-%m-%d %H:%M:%S.%f")[:-3]

    def _mockdate_relative(self, base_time=None):
        if base_time is None:
            base_time = self._current_time or datetime.now()

        if not self._rule:
            return base_time.strftime("%Y-%m-%d %H:%M:%S")

        rule_str = self._rule.strip()

        # ✅ 支持纯相对偏移：+300, -1800
        if rule_str.startswith(("+", "-")) and rule_str[1:].isdigit():
            try:
                seconds = int(rule_str)
                new_time = base_time + timedelta(seconds=seconds)
                return new_time.strftime("%Y-%m-%d %H:%M:%S")
            except:
                pass

        # 否则当作绝对时间格式（用 exrex 生成）
        return self._mockdate()

    def _format(self):
        """
        根据 format@date@... 中的 pattern 格式化当前时间上下文。
        支持：YYYY, MM, dd, HH, mm, ss, SSS
        """

        dt = self._current_time or datetime.now()
        rule_content = self._rule.strip()
        # 如果是 date@... 形式，提取后面的格式
        if rule_content.startswith("date@"):
            fmt_pattern = rule_content[5:]  # 跳过 "date@"
        else:
            # 直接当作格式模板
            fmt_pattern = rule_content
        if "SSS" in fmt_pattern:
            ms = dt.microsecond // 1000
            fmt_str = fmt_str.replace("SSS", f"{ms:03d}")

        # 映射其他格式符
        replacements = {
            "YYYY": dt.strftime("%Y"),
            "YY": dt.strftime("%y"),
            "MM": dt.strftime("%m"),
            "dd": dt.strftime("%d"),
            "HH": dt.strftime("%H"),
            "mm": dt.strftime("%M"),
            "ss": dt.strftime("%S"),
        }

        for token, value in replacements.items():
            if token in fmt_pattern:
                fmt_pattern = fmt_pattern.replace(token, value)

        return fmt_pattern


# 注册对象
ruleAnalysis = RuleAnalysis()


def rule(rule_reg, ruleAnalysis: RuleAnalysis = None) -> None:
    """
    解析规则并模拟生成数据。

    参数:
    - rule_reg (str): 包含数据类型和规则的字符串，格式为 "数据类型@规则"。
    - ruleAnalysis (RuleAnalysis)：解析器对象

    支持的数据类型:
    - str: 字符串类型
    - date: 日期类型
    - int: 整数类型
    - none: 原文输出

    支持合并:
    通过&&符号 多项合并 形如 str@xxx&&date@yyyy-MM&&int@4

    str支持的规则:
    - 使用正则，输出一个完全匹配的字符串
    - 支持输入 yyyy 输出2024-10
    - 支持输入 MM 输出2024-10
    - 支持输入 dd 输出2024-10
    - 支持输入 HH 输出11
    - 支持输入 mm 输出24
    - 支持输入 ss 输出56
    - 支持输入任意拼接的上述格式 yyyy-MM-dd HH:mm:ss 输出2024-10-25 11:24:56
    - 支持生成唯一编号 uid uid(20) 输出20位随机串 如果没有数值，默认为20位长度 注意在应用括号的时候，进行转义
    - 支持生成IPV4地址 ipv4
    - 支持生成IPV6地址 ipv6

    date支持的规则
    - 允许输入日期正则
    - 支持输入 yyyy 输出2024
    - 支持输入 MM 输出10
    - 支持输入 dd 输出24
    - 支持输入 HH 输出11
    - 支持输入 mm 输出24
    - 支持输入 ss 输出56
    - 支持输入任意拼接的上述格式 yyyy-MM-dd HH:mm:ss 输出2024-10-25 11:24:56


    num支持的规则
    - 输入数值，用于控制字符串长度
    - 输入非数值，返回长度为1的数值
    - 输入数值正则，返回随机数值 支持 Number 格式

    choice支持的规则:
    - choice@file.json          # 读取 JSON 所有叶子节点值
    - choice@file.json!data[*].name  # 读取 JSON 指定路径
    - choice@data.xlsx          # 读取 Excel 第一个 Sheet 第一列
    - choice@data.xlsx!部门     # 读取指定 Sheet 的"部门"列（默认第一个 Sheet）
    - choice@Sheet2!城市        # 读取 "Sheet2" 的"城市"列

    choice_row@file.xlsx!col → 选中整行并缓存（按文件）
    choice_ref@字段 → 引用最近文件缓存行中的字段
    choice_ref@file.xlsx!字段 → 明确引用指定文件缓存行中的字段

    返回:
    生成的数据，根据规则解析而来。
    """

    if not rule_reg:
        return ""

    pattern = r"(?P<str>str\s*@\s*[^&&]+)|(?P<date>date\s*@\s*[^&&]+)|(?P<num>num\s*@\s*[^&&]+)|(?P<none>none\s*@\s*[^&&]+)|(?P<format>format\s*@\s*[^&&]+)"

    # 查找所有匹配
    matches = list(re.finditer(pattern, rule_reg))

    # 判断是否有匹配内容
    if len(matches) < 1:
        return rule_reg  # 如果没有匹配，则返回原始字符串

    # 用于存储匹配结果
    results = []

    # 遍历所有匹配项
    for match in matches:
        # 获取匹配的分组
        for group_name, group_value in match.groupdict().items():
            if group_value:  # 只处理非空分组
                # 去除空白字符
                # if group_name != "none":
                #     group_value = re.sub(r"\s+", "", group_value)
                results.append(group_value)

    exrex_str = ""

    if ruleAnalysis is None:
        ruleAnalysis = RuleAnalysis()

    for sub_rule_reg in results:

        data_type, rule_str = sub_rule_reg.split("@", 1)

        rule_func = {
            "str": lambda rule: (ruleAnalysis._setrule(rule)._mockstr()),
            "date": lambda r: ruleAnalysis._setrule(r)._mockdate_relative(
                ruleAnalysis._current_time
            ),
            "num": lambda rule: (ruleAnalysis._setrule(rule)._mocknum()),
            "none": lambda rule: rule,
            "format": lambda rule: (ruleAnalysis._setrule(rule)._format()),
        }

        keys = rule_func.keys()
        if any(key in data_type for key in keys) == False:
            exrex_str += sub_rule_reg
            continue

        try:
            exrex_str += rule_func[data_type](rule_str)
        except Exception as e:
            print(f"规则处理错误: {sub_rule_reg}, 错误: {str(e)}")
            exrex_str += sub_rule_reg

    return exrex_str if exrex_str != "" else rule_reg


def apply(original_value: str, rule_str: str, ruleAnalysis: RuleAnalysis = None) -> str:
    """
    根据原始值和规则字符串，生成新的字符串值。
    规则中可通过 {value} 引用原始值。

    示例:
        apply_rule("U1001", "str@PRE_{value}_uid(6)") → "PRE_U1001_XK9Q2M"
        apply_rule("2025-04-01 10:00:00", "date@{value}+300") → "2025-04-01 10:05:00"
        apply_rule("北京", "str@{value}&&str@_Office") → "北京_Office"
    """
    if ruleAnalysis is None:
        ruleAnalysis = RuleAnalysis()

    # 替换 {value} 为原始值
    processed_rule = rule_str.replace("{value}", original_value)

    # 如果规则中包含时间相关，尝试解析 original_value 为 datetime 设置上下文
    # ✅ 改进：只要规则中包含 format@date@ 或 date@，就尝试解析 original_value 为时间
    if "format@date@" in processed_rule or "date@" in processed_rule:
        try:
            from datetime import datetime

            for fmt in [
                "%Y-%m-%d %H:%M:%S",
                "%Y-%m-%d %H:%M:%S.%f",
                "%Y-%m-%d",
                "%Y/%m/%d %H:%M:%S",
                "%Y/%m/%d",
            ]:
                try:
                    dt = datetime.strptime(original_value.strip(), fmt)
                    ruleAnalysis.set_current_time(dt)
                    break
                except ValueError:
                    continue
        except Exception:
            pass  # 解析失败则忽略

    try:
        result = rule(processed_rule, ruleAnalysis)
        return str(result) if result is not None else ""
    except Exception as e:
        return f"ERROR: {str(e)}"


def apply_rule(
    original_value: str, rule_str: str, ruleAnalysis: RuleAnalysis = None
) -> str:
    """
    根据 rule_str 对 original_value 进行二次计算。
    支持嵌套规则，如: apply@date@+ (num@\\d{1,2})
    先解析所有 (type@rule) 形式的嵌套规则，再拼接执行。
    """
    if rule_str.startswith("apply@"):
        rule_str = rule_str[6:]  # 去掉 "apply@"

    if ruleAnalysis is None:
        ruleAnalysis = RuleAnalysis()

    # 定义嵌套规则的正则：匹配 (str@...), (date@...), (num@...), (choice@...), (format@date...) 等
    # 注意：括号必须是英文括号，且内部不能嵌套括号（简单场景）
    nested_pattern = (
        r"\((?P<type>str|date|num|choice(?:_row)?|format)@(?P<content>[^)]+)\)"
    )

    def replace_nested(match):
        nested_rule = f"{match.group('type')}@{match.group('content')}"
        # 递归解析这个嵌套规则（不依赖 original_value）
        try:
            result = rule(nested_rule, ruleAnalysis)
            return str(result)
        except Exception as e:
            return f"ERROR_IN_NESTED:({nested_rule})"

    # 循环替换，直到没有嵌套（处理多层嵌套）
    # 例如: (str@(num@5)) → 先替换成 (str@5) → 再替换成 5
    while re.search(nested_pattern, rule_str):
        rule_str = re.sub(nested_pattern, replace_nested, rule_str)

    # 现在 rule_str 已无嵌套，交给 apply 执行
    return apply(original_value, rule_str, ruleAnalysis)



def extra_rule(
    extra_data: Dict[str, List[Dict[str, Any]]], rule_str: str
) -> Union[str, List[str]]:
    """
    根据 rule_str 从 extra_data 中提取并处理数据。

    Args:
        extra_data: 外部数据，格式 {filename: [ {field: value, ...}, ... ]}
        rule_str: 规则字符串，格式如 "extra@ASSET.xlsx@ASSET_ID@join@,"

    Returns:
        str: 如果操作是 join，返回连接后的字符串
        str: 如果操作是 length，返回字段值数量的字符串
        List[str]: 如果操作是 each，返回字段值列表
        List[str]: 如果是 select[field], 返回字段过滤出来的集合extra@ASSET.xlsx@select[ASSET_ID]@1212,
    """
    parts = rule_str.split("@")

    if not parts or parts[0] != "extra":
        raise ValueError(f"Invalid rule format: {rule_str}. Must start with 'extra@'")

    if len(parts) < 4:
        raise ValueError(f"Rule too short: {rule_str}. Expected at least 4 parts.")

    _, data_key, field_name, op = parts[0], parts[1], parts[2], parts[3]
    op_param = parts[4] if len(parts) > 4 else ","  # 默认分隔符为逗号

    # 1. 获取外部数据列表
    records = extra_data.get(data_key, [])
    match = re.search(r"select\[(.*?)\]", op)
    if not records:
        # 若无数据，根据操作返回空值
        if op == "each":
            return []
        elif match:
            return []
        elif op == "length":
            return "0"
        else:
            return ""  # join 返回空字符串

    # 2. 提取指定字段的值（安全：缺失字段视为 None 或空字符串）
    values = []
    for record in records:
        if match:
            select_field = match.group(1).strip()
            if not record.get(select_field) == op_param:
                continue

        val = record.get(field_name)
        if val is None:
            val = ""  # 可根据需求改为跳过或保留 None
        else:
            val = str(val)
        values.append(val)

    # 3. 根据操作类型处理
    if op == "join":
        return op_param.join(values)
    elif op == "each":
        return values
    elif match:
        return values
    elif op == "length":
        return str(len(values))
    else:
        raise ValueError(f"Unsupported operation: {op}. Use 'join' or 'each'.")


__all__ = ["ruleAnalysis", "rule", "apply_rule", "extra_rule"]


if __name__ == "__main__":
    # 创建测试文件
    test_data1 = {
        "user_id": ["U001", "U002", "U003"],
        "user_name": ["张三", "李四", "王五"],
        "department": ["技术部", "市场部", "财务部"],
    }

    test_data2 = {
        "order_id": ["O1001", "O1002", "O1003"],
        "product": ["手机", "电脑", "平板"],
        "price": [2999, 5999, 3999],
    }

    # 创建测试Excel文件
    df1 = pd.DataFrame(test_data1)
    df1.to_excel("users.xlsx", index=False)

    df2 = pd.DataFrame(test_data2)
    df2.to_excel("orders.xlsx", index=False)

    print("=== 测试多文件 choice_row + choice_ref ===")
    print("1. 选中 users.xlsx 中的一行:")
    print(rule("str@choice_row@users\\.xlsx!user_id", ruleAnalysis))
    print("   引用 user_name:", rule("str@choice_ref@user_name", ruleAnalysis))
    print("   引用 department:", rule("str@choice_ref@department", ruleAnalysis))

    print("\n2. 选中 orders.xlsx 中的一行:")
    print(rule("str@choice_row@orders\\.xlsx!order_id", ruleAnalysis))
    print("   引用 product:", rule("str@choice_ref@product", ruleAnalysis))
    print("   引用 price:", rule("str@choice_ref@price", ruleAnalysis))

    print("\n3. 显式跨文件引用 users.xlsx 的 user_name:")
    print(rule("str@choice_ref@users\\.xlsx!user_name", ruleAnalysis))

    print("\n4. 测试静态 choice 数组:")
    print(rule('str@choice@\\["研发部", "市场部", "人事部"\\]', ruleAnalysis))

    print("\n5. 测试 UID 和 NUM:")
    print(rule("str@Mainuid\\(16\\)_uid\\(3\\)", ruleAnalysis))
    print(rule("num@\\d{1,2}", ruleAnalysis))
    print(rule("str@[01]{1}", ruleAnalysis))

    print("\n6. 字符串拼接 + uid:")
    # 示例1: 字符串拼接 + uid
    print(apply("U1001", "str@ORDER_{value}_uid(6)", ruleAnalysis))
    # 输出: ORDER_U1001_A7F3K9

    # 示例2: 时间偏移
    print(apply("2025-04-01 10:00:00", "date@+300", ruleAnalysis))
    # 输出: 2025-04-01 10:05:00

    # 示例3: 多规则合并
    print(apply("北京", "str@{value}&&str@_Branch", ruleAnalysis))
    # 输出: 北京_Branch

    # 示例4: choice 引用（假设已加载文件）
    print(
        apply(
            "技术部",
            "str@choice_ref@users\\.xlsx!department&&str@_&&str@choice_ref@users\\.xlsx!user_name",
            ruleAnalysis,
        )
    )
    # 输出: 技术部_张经理

     # 示例5: apply 应用
    print(apply_rule("2025-04-01 10:00:00", "apply@date@+(num@60~300)", ruleAnalysis))

    print(
        apply_rule(
            "2025-04-01 10:00:00",
            "apply@str@测试任务(format@date@YYYYMMddHHmmss)",
            ruleAnalysis,
        )
    )


    # 清理测试文件
    for f in ["users.xlsx", "orders.xlsx"]:
        if os.path.exists(f):
            os.remove(f)
